﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace LowLevelGraphics.Filter
{
    /// <summary>
    /// Full Multigrid (FMG) Algorithm implementation for solution of linear elliptic equation.
    /// </summary>
    public sealed class MultigridSolver
    {
        private const int NPRE = 1; // number of relaxation sweeps before...
        private const int NPOST = 1; // ...and after the coarse-grid correction is computed
        private const int NGMAX = 15;

        /// <summary>
        /// Get nearest larger dimension length that is acceptable by the algorithm.
        /// </summary>
        public int getSize(int width, int height)
        {

            return (int)(Math.Pow(2, Math.Ceiling(Math.Log(Math.Max(width, height), 2))) + 2);
        }

        /// <summary>
        /// Solve linear elliptic equation.
        /// </summary>
        /// <param name="u">Right-hand side (input).</param>
        /// <param name="n">Dimension (must be of the form 2^j + 1, j = 0,1,...).</param>
        /// <param name="nCycle">Number of V cycles at each level.</param>
        public void mglin(double[,] u, int n, int ncycle)
        {

            n--; // a little correction because for the different indexing scheme

            int j, jcycle, jj, jpost, jpre, nf, ngrid, nn;
            double[][,] ires = new double[NGMAX + 1][,];
            double[][,] irho = new double[NGMAX + 1][,];
            double[][,] irhs = new double[NGMAX + 1][,];
            double[][,] iu = new double[NGMAX + 1][,];

            int ng = 0;

            nn = n;

            while ((nn >>= 1) != 0)
                ng++;

            if (n != (1 + (1 << ng)))
                return;

            if (ng > NGMAX)
                return;

            nn = (n / 2 + 1);
            ngrid = (ng - 1);

            // allocate storage form right-hand side on grid ng-1 and fill it by restricting from
            // the fine grid. Similarly allocate storage and fill right-hand side on all coarse
            // grids.
            irho[ngrid] = new double[nn + 1, nn + 1];

            rstrct(irho[ngrid], u, nn);

            while (nn > 3)
            {

                nn = (nn / 2 + 1);

                irho[--ngrid] = new double[nn + 1, nn + 1];

                rstrct(irho[ngrid], irho[ngrid + 1], nn);
            }

            nn = 3;

            iu[1] = new double[nn + 1, nn + 1];
            irhs[1] = new double[nn + 1, nn + 1];

            // initial solution on coarsest grid
            slvsml(iu[1], irho[1]);

            irho[1] = null;

            ngrid = ng;

            // nested iteration loop
            for (j = 2; j <= ngrid; j++)
            {

                nn = (2 * nn - 1);

                iu[j] = new double[nn + 1, nn + 1];
                irhs[j] = new double[nn + 1, nn + 1];
                ires[j] = new double[nn + 1, nn + 1];

                interp(iu[j], iu[j - 1], nn);

                // interpolate from coarse grid to next finer grid
                copy(irhs[j], (j != ngrid) ? irho[j] : u, nn); // set up right-hand side

                // V-cycle loop
                for (jcycle = 1; jcycle <= ncycle; jcycle++)
                {

                    nf = nn;

                    // downward stroke of the V.
                    for (jj = j; jj >= 2; jj--)
                    {

                        // pre-smoothing
                        for (jpre = 1; jpre <= NPRE; jpre++)
                            relax(iu[jj], irhs[jj], nf);

                        resid(ires[jj], iu[jj], irhs[jj], nf);

                        nf = (nf / 2 + 1);

                        rstrct(irhs[jj - 1], ires[jj], nf);

                        // restriction of the residual is the next right-hand side
                        fill0(iu[jj - 1], nf); // zero for initial guess in next relaxation
                    }

                    // bottom of V: solve on coarsest grid
                    slvsml(iu[1], irhs[1]);

                    nf = 3;

                    //upward stroke of V.
                    for (jj = 2; jj <= j; jj++)
                    {

                        nf = (2 * nf - 1);

                        addint(iu[jj], iu[jj - 1], ires[jj], nf);

                        // use res for temporary storage inside addint
                        for (jpost = 1; jpost <= NPOST; jpost++) // post-smoothing
                            relax(iu[jj], irhs[jj], nf);
                    }
                }
            }

            // return solution in u
            copy(u, iu[ngrid], n);

            for (nn = n, j = ng; j >= 2; j--, nn = (nn / 2 + 1))
            {

                ires[j] = null;
                irhs[j] = null;
                iu[j] = null;

                if (j != ng)
                    irho[j] = null;
            }

            irhs[1] = null;
            iu[1] = null;
        }

        /// <summary>
        /// Half-weighting restriction.
        /// </summary>
        /// <param name="uc">Coarse-grid solution (output).</param>
        /// <param name="uf">Fine-grid solution (input).</param>
        /// <param name="nc">Coarse-grid dimension.</param>
        private void rstrct(double[,] uc, double[,] uf, int nc)
        {

            int ic, iif, jc, jf;
            int ncc = (2 * nc - 1);

            // interior points
            for (jf = 3, jc = 2; jc < nc; jc++, jf += 2)
                for (iif = 3, ic = 2; ic < nc; ic++, iif += 2)
                    uc[ic, jc] = (0.5 * uf[iif, jf] + 0.125 * (uf[iif + 1, jf] + uf[iif - 1, jf] + uf[iif, jf + 1] + uf[iif, jf - 1]));

            // boundary points
            for (jc = 1, ic = 1; ic <= nc; ic++, jc += 2)
            {

                uc[ic, 1] = uf[jc, 1];
                uc[ic, nc] = uf[jc, ncc];
            }

            for (jc = 1, ic = 1; ic <= nc; ic++, jc += 2)
            {

                uc[1, ic] = uf[1, jc];
                uc[nc, ic] = uf[ncc, jc];
            }
        }

        /// <summary>
        /// Coarse-to-fine prolongation by bilinear interpolation.
        /// </summary>
        /// <param name="uf">Fine-grid solution (output).</param>
        /// <param name="uc">Coarse-grid solution (input).</param>
        /// <param name="nf">Fine-grid dimension.</param>
        private void interp(double[,] uf, double[,] uc, int nf)
        {

            int ic, iif, jc, jf, nc;

            nc = (nf / 2 + 1);

            for (jc = 1, jf = 1; jc <= nc; jc++, jf += 2) // do elements that are copies
                for (ic = 1; ic <= nc; ic++)
                    uf[2 * ic - 1, jf] = uc[ic, jc];

            for (jf = 1; jf <= nf; jf += 2) // do odd-numbered columns, interpolating vertically
                for (iif = 2; iif < nf; iif += 2)
                    uf[iif, jf] = (0.5 * (uf[iif + 1, jf] + uf[iif - 1, jf]));

            for (jf = 2; jf < nf; jf += 2) // do even-numbered columns, interpolating horizontally
                for (iif = 1; iif <= nf; iif++)
                    uf[iif, jf] = (0.5 * (uf[iif, jf + 1] + uf[iif, jf - 1]));
        }

        /// <summary>
        /// Does coarse-to-fine interpolation and adds result to uf.
        /// </summary>
        /// <param name="uf">Fine-grid solution (output).</param>
        /// <param name="uc">Coarse-grid solution (input).</param>
        /// <param name="res">Temporary storage.</param>
        /// <param name="nf">Fine-grid dimension.</param>
        private void addint(double[,] uf, double[,] uc, double[,] res, int nf)
        {

            int i, j;

            interp(res, uc, nf);

            for (j = 1; j <= nf; j++)
                for (i = 1; i <= nf; i++)
                    uf[i, j] += res[i, j];
        }

        /// <summary>
        /// Solution to the model problem on the coarsest grid, where h = 1/2.
        /// </summary>
        /// <param name="u">Solution (output).</param>
        /// <param name="rhs">Right-hand side function (input).</param>
        private void slvsml(double[,] u, double[,] rhs)
        {

            double h = 0.5;

            fill0(u, 3);

            u[2, 2] = (-h * h * rhs[2, 2] / 4.0);
        }

        /// <summary>
        /// Red-black Gauss-Seidel relaxation for model problem.
        /// </summary>
        /// <param name="u">Current value of the solution to be updated (input/output).</param>
        /// <param name="rhs">Right-hand side function.</param>
        private void relax(double[,] u, double[,] rhs, int n)
        {

            int i, ipass, isw, j;
            double h, h2;

            int jsw = 1;

            h = (1.0 / (n - 1));
            h2 = (h * h);

            for (ipass = 1; ipass <= 2; ipass++, jsw = (3 - jsw))
            { // red and black sweeps

                isw = jsw;

                for (j = 2; j < n; j++, isw = (3 - isw))
                    for (i = (isw + 1); i < n; i += 2) // Gauss-Seidel formula.
                        u[i, j] = (0.25 * (u[i + 1, j] + u[i - 1, j] + u[i, j + 1] + u[i, j - 1] - h2 * rhs[i, j]));
            }
        }

        /// <summary>
        /// Returns minus the residual for the model problem.
        /// </summary>
        /// <param name="res">Return value.</param>
        /// <param name="u">Input quantity.</param>
        /// <param name="rhs">Input quantity.</param>
        private void resid(double[,] res, double[,] u, double[,] rhs, int n)
        {

            int i, j;
            double h, h2i;

            h = (1.0 / (n - 1));
            h2i = (1.0 / (h * h));

            // interior points
            for (j = 2; j < n; j++)
                for (i = 2; i < n; i++)
                    res[i, j] = (-h2i * (u[i + 1, j] + u[i - 1, j] + u[i, j + 1] + u[i, j - 1] - 4.0 * u[i, j]) + rhs[i, j]);

            // boundary points
            for (i = 1; i <= n; i++)
                res[i, 1] = res[i, n] = res[1, i] = res[n, i] = 0.0;
        }

        /// <summary>
        /// Copies ain to aout.
        /// </summary>
        private void copy(double[,] aout, double[,] ain, int n)
        {

            int i, j;

            for (i = 1; i <= n; i++)
                for (j = 1; j <= n; j++)
                    aout[j, i] = ain[j, i];
        }

        /// <summary>
        /// Fills u with zeros.
        /// </summary>
        private void fill0(double[,] u, int n)
        {

            int i, j;

            for (j = 1; j <= n; j++)
                for (i = 1; i <= n; i++)
                    u[i, j] = 0.0;
        }
    }
}
